import cv2
import numpy as np
import os
from PIL import Image, ImageTk
from tkinter import messagebox
import tkinter as tk
import pandas as pd
import openpyxl


class gui(tk.Tk):
    def __init__(self):
        super().__init__()
        self.Index = -1
        self.Index1 = -1
        self.entry1 = None
        self.name = None
        self.allPic = tk.Image
        self.length = 0
        self.img = 0
        self.mark = []
        self.number = []
        self.data = {}
        self.rightAns = np.array([[0, 0, 1, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 0, 1]])
        self.title("答题卡识别系统by侯智鹏")
        self.geometry("900x700")

    def createButton(self):
        lb1 = tk.Label(self, text='欢迎来到答题卡阅卷系统', fg='blue', font=("宋体", 30))
        lb1.place(relx=0.1, rely=0.02, width=900 * 0.8)
        lb2 = tk.Label(self, text="文件路径：", fg='black', font=("宋体", 15))
        lb2.place(x=10, rely=0.85, width=100)
        lb3 = tk.Label(self, text="批改前", fg='red', font=("宋体", 20))
        lb3.place(relx=0.2, rely=0.1, width=80)
        lb4 = tk.Label(self, text="批改后", fg='red', font=("宋体", 20))
        lb4.place(relx=0.7, rely=0.1, width=80)
        lb5 = tk.Label(self, text="答案为：CAABE", fg='blue', font=("宋体", 20))
        lb5.place(rely=0.92, relx=0.4, width=2 * 450 * 0.2)
        entry1 = tk.Entry(self)
        entry1.place(x=111, rely=0.85, width=130)
        self.entry1 = entry1
        button1 = tk.Button(self, text="开始阅卷", font=("宋体", 15), command=self.effects)
        button1.place(x=400, rely=0.85, width=100)
        button2 = tk.Button(self, text="下一张", font=("宋体", 15), command=self.loadPicture)
        button2.place(x=550, rely=0.85, width=100)
        button3 = tk.Button(self, text="存入excel", font=("宋体", 15), command=self.mark2excel)
        button3.place(x=700, rely=0.85, width=100)
        button4 = tk.Button(self, text="读取文件", font=("宋体", 15), command=self.judge)
        button4.place(x=250, rely=0.85, width=80)

    def judge(self):
        a = self.entry1.get()
        path = a + '/'
        allName = []
        try:
            for name in os.listdir(path):
                if name.endswith('.png'):
                    wholeName = path + name
                    allName.append(wholeName)
            if len(allName) != 0:
                self.Index1 += 1
                if self.Index1 > 0:
                    messagebox.showinfo("文件读取", "已经成功读取，请勿重复读取")
                else:
                    self.length = len(allName)
                    messagebox.showinfo("文件读取", "读取成功")
        except:
            messagebox.showinfo("文件读取", "该文件无内容或不存在该文件")

    def effects(self):
        try:
            thresh, Gau = solvePicture(self.img, (3, 3), 0, way=cv2.THRESH_OTSU | cv2.THRESH_BINARY)
            allPts = findShape(thresh)
            wrap = fourPointsTrans(Gau, allPts)
            beforePic, mark = findUserAns(wrap, self.rightAns)
            self.mark.append(mark)
            self.number.append(self.name[7:-4])
            beforePic = Image.fromarray(cv2.cvtColor(beforePic, cv2.COLOR_BGR2RGB))
            beforePic = ImageTk.PhotoImage(beforePic)
            beforePicMain = Image.open(self.name)
            beforePicMain = ImageTk.PhotoImage(beforePicMain)
            lb3 = tk.Label(self, image=beforePicMain)
            lb3.place(x=20, rely=0.15, width=400, relheight=0.6)
            lb4 = tk.Label(self, image=beforePic)
            self.allPic = beforePic, beforePicMain
            lb4.place(x=480, rely=0.15, width=400, relheight=0.6)
        except:
            messagebox.showinfo("提示", "请输入文件夹正确名字或者先按下一张再开始改卷")

    def loadPicture(self):
        a = self.entry1.get()
        path = a + '/'
        allName = []
        try:
            for name in os.listdir(path):
                if name.endswith('.png'):
                    wholeName = path + name
                    allName.append(wholeName)
            self.Index += 1
            if self.Index >= self.length:
                if self.length == 0:
                    messagebox.showinfo("文件读取", "请输入文件夹正确名字")
                    self.Index = -1
                else:
                    messagebox.showinfo("文件读取", "已经是最后一张了")
            else:
                self.name = allName[self.Index]
                img = cv2.imread(self.name, cv2.IMREAD_COLOR)
                self.img = img
                beforePicMain = Image.open(self.name)
                beforePicMain = ImageTk.PhotoImage(beforePicMain)
                lb3 = tk.Label(self, image=beforePicMain)
                self.allPic = beforePicMain
                lb3.place(x=20, rely=0.15, width=400, relheight=0.6)
        except:
            pass

    def mark2excel(self):
        if self.length == 0:
            messagebox.showinfo("excel读入", "无法读入excel，请输入文件夹正确名字或确保文件夹内有内容")
        else:
            self.data['学号'] = self.number
            self.data['分数'] = self.mark
            writer = pd.ExcelWriter("../mark/mark.xlsx")
            data = pd.DataFrame(self.data)
            data.to_excel(writer)
            writer.save()
            messagebox.showinfo("excel读入", "excel读入成功")


def solvePicture(img, GauKernel, threshold, way=cv2.THRESH_BINARY):
    Gau = cv2.GaussianBlur(img, GauKernel, sigmaX=0, sigmaY=0)
    grey = cv2.cvtColor(Gau, cv2.COLOR_BGR2GRAY)
    ret, thresh = cv2.threshold(grey, threshold, 255, way)
    # 返回二值图和灰度图
    return thresh, Gau


def findShape(thresh):
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    allPts = []
    for index, context in enumerate(contours):
        (x, y, w, h) = cv2.boundingRect(context)
        if w > 100 and h > 100:
            allPts.append((x, y, w, h))
    allPts = sorted(allPts, key=lambda r: r[0])
    return allPts


def fourPointsTrans(img, pts):
    subPtsFeature = pts[0]
    (x, y, w, h) = subPtsFeature
    rect = np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h]], dtype=np.float32)
    (ls, rs, rx, lx) = rect
    # 重新计算新的图像的高度和宽度
    width1 = np.sqrt(np.sum(rx[0] - lx[0]) ** 2)
    width2 = np.sqrt(np.sum(rs[0] - ls[0]) ** 2)
    maxWidth = max(int(width1), int(width2))
    height1 = np.sqrt(np.sum(rs[1] - rx[1]) ** 2)
    height2 = np.sqrt(np.sum(ls[1] - lx[1]) ** 2)
    maxHeight = max(int(height1), int(height2))
    dst = np.array([[0, 0], [maxWidth - 1, 0], [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]], dtype=np.float32)
    M = cv2.getPerspectiveTransform(rect, dst)
    wrap = cv2.warpPerspective(img, M, (maxWidth, maxHeight))
    return wrap


def findUserAns(img, rightAnswer):
    imgCopy = img
    thresh, Gau = solvePicture(img, (3, 3), 0, way=cv2.THRESH_OTSU | cv2.THRESH_BINARY_INV)
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    ans = []
    for i in contours:
        area = float(cv2.contourArea(i))
        length = float(cv2.arcLength(i, True))
        if area <= 0 or length <= 0:
            continue
        if 7.05 < area / length < 10.5:
            ans.append(i)
    circle = []
    for i in ans:
        (x, y), r = cv2.minEnclosingCircle(i)
        center = (int(x), int(y))
        r = int(r)
        circle.append((center, r))
    # 按照y坐标从小到大排序
    circle.sort(key=lambda t: t[0][1])
    A = []
    for i in range(1, 6):
        now = circle[(i - 1) * 5: i * 5]
        now.sort(key=lambda p: p[0][0])
        A.extend(now)
    # 对于A中的每个选项圆,计算它所有覆盖的坐标，然后判断这些坐标在二值图对应的值，统计白色点个数，占比高的话被选中
    da = []
    for i in A:
        dots = []
        # 筛选在圆中的点
        for m in range(-i[1], i[1] + 1):
            for n in range(-i[1], i[1] + 1):
                dot2 = (i[0][0] + m, i[0][1] + n)
                result = ((i[0][0] - dot2[0]) ** 2 + (i[0][1] - dot2[1]) ** 2) ** 0.5
                if result <= i[1]:
                    dots.append(dot2)
        allDots = len(dots)
        whiteDots = 0
        # 计算白色点
        for j in dots:
            if thresh[j[1]][j[0]] == 255:
                whiteDots += 1
        if whiteDots / allDots >= 0.55:
            da.append(1)
        else:
            da.append(0)
    da = np.array(da).reshape((5, 5))
    connectAns = da + rightAnswer
    m, n = connectAns.shape
    correctIndex = []
    for i in range(m):
        for j in range(n):
            if connectAns[i][j] == 2:
                correctIndex.append((i, j))
    correctLength = len(correctIndex)
    mark = int((correctLength / 5.0) * 100)
    for i in range(correctLength):
        subIndex = correctIndex[i]
        (correctX, correctY), correctRad = A[subIndex[0] * 5 + subIndex[1]]
        cv2.circle(imgCopy, (correctX, correctY), correctRad, (0, 255, 0), 3)
    wrongAnsIndex = []
    for i in range(m):
        for j in range(n):
            if connectAns[i][j] != 2 and rightAnswer[i][j] == 1:
                wrongAnsIndex.append((i, j))
    wrongLength = len(wrongAnsIndex)
    for i in range(wrongLength):
        subIndex = wrongAnsIndex[i]
        (wrongX, wrongY), wrongRad = A[subIndex[0] * 5 + subIndex[1]]
        cv2.circle(imgCopy, (wrongX, wrongY), wrongRad, (0, 0, 255), 3)
    font = "mark=%d" % mark
    cv2.putText(imgCopy, font, (25, 412), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
    return imgCopy, mark


def testMain():
    win = gui()
    win.createButton()
    win.mainloop()


if __name__ == '__main__':
    testMain()
